#include <iomanip>
#include <exception>
#include <cmath>
#include <lapacke.h>
#include "Embedded/PartialDer.h"
#include "Embedded/SourceStencil.impl.h"
#include "Embedded/EvaluateMonomials.h"
#include "Embedded/ConditionNumber.h"

template <int Dim, int pD, int _fitOrder>
PartialDer<Dim, pD, _fitOrder>::PartialDer(const EmGrid &aGrid)
    : emGrid(aGrid)
{
  // 0.
  struct _partialDerTransform : public partialDerTransform { enum { fitOrder = _fitOrder }; };
  // 1. Construct the stencils.
  SourceStencil<Dim> ss(emGrid, _partialDerTransform(), KBox_Diffusion_Dominated<Dim>(emGrid, _fitOrder / 2));
  // 2. Find the approximation.
  using NodeType = typename EmGrid::NodeType;
  const auto &nodeLabels = emGrid.nodeLabels;
  const auto &bdryPoints = emGrid.bdryPoints;
  const auto dx = emGrid.spacing();
  const int n = fitOrder - 1;
  const int nBasis = binom(n+Dim, Dim);
  for(const auto &stencil : ss.vStencil) {
    const auto destIdx = stencil.destIdx;
    const auto refLoc  = (destIdx + 0.5) * dx;
    // 2.1. calculate the sample matrix.
    Tensor<Real, 2> sampleMatrix(Vec<int, 2> {nBasis, nBasis});
    sampleMatrix = 0.0;
    Real *where2Fill = sampleMatrix.data();
    for(const auto &jdx : stencil.locations) {
      rVec p;
      if(nodeLabels(jdx) == NodeType::Interior) {
        p = (jdx + 0.5) * dx;
      } else if(nodeLabels(jdx) == NodeType::Ghost) {
        // use the nearest boudary point as boundary condition
        p = bdryPoints(jdx);
      } else {
        assert(0);
      }
      EvaluateMonomials<Dim,n>::doEvaluate(identityOfMonomials(),
                                           [](const Real &x, const auto &) { return x; },
                                           where2Fill, p-refLoc, rVec(0.0), dx);
    }
    assert((where2Fill - sampleMatrix.data()) == sampleMatrix.volume());
#ifndef NDEBUG
    const Real cn = estimateSquareCond(sampleMatrix);
    if(cn >= dangerousCN || _dbglevel >= 2)
      dbgcout1 << "The linear system for " << destIdx
               << " has condition number " << std::setprecision(2) << std::scientific << cn
               << ". \n" << std::defaultfloat;
    dbgcout3 << std::showpos << std::setprecision(3) << sampleMatrix << "\n";
    dbgcout3 << std::defaultfloat;
#endif
    // 2.2. calculate the RHS.
    Tensor<Real, 1> sampleRHS(nBasis);
    sampleRHS = 0.0;
    where2Fill = sampleRHS.data();
    EvaluateMonomials<Dim,n>::doEvaluate(_partialDerTransform(),
                                         [](const Real &x, const auto &) { return x; },
                                         where2Fill, rVec(0.0)/*refLoc-refLoc*/, rVec(0.0), dx);
    // 2.3. solve the linear system to obtain the coefficients.
    Tensor<int, 1> ipiv(nBasis);
    auto info = LAPACKE_dgesv(LAPACK_COL_MAJOR, nBasis, 1, sampleMatrix.data(), nBasis, ipiv.data(), sampleRHS.data(), nBasis);
    if(info != 0)
      throw std::runtime_error("(PartialDer::PartialDer) DGESV failed.");
    // 2.4. Save the coefficients in sparse matrix format.
    matOp.insertRow(destIdx, nBasis, stencil.locations.data(), sampleRHS.data());
    dbgcout2 << sampleRHS << "\n";
  }
}

template <int Dim, int pD, int _fitOrder>
void PartialDer<Dim, pD, _fitOrder>::apply(const Tensor<Real, Dim> &aData, Tensor<Real, Dim> &outData) const
{
  // 1. apply the standard scheme.
  auto dx = emGrid.spacing();
  auto u = iVec::unit(pD);
  if(_fitOrder == 5) {
    ddfor(emGrid, [&](const iVec &idx) {
      outData(idx) =
          (-aData(idx + u * 2) + aData(idx + u) * 8 - aData(idx - u) * 8 + aData(idx - u * 2)) / (12.0 * dx[pD]);
    });
  } else if(_fitOrder == 7) {
    ddfor(emGrid, [&](const iVec &idx) {
      outData(idx) = (aData(idx+u*3) - aData(idx+u*2) * 9 + aData(idx+u) * 45
          - aData(idx-u) * 45 + aData(idx-u*2) * 9 - aData(idx-u*3)) / (60.0 * dx[pD]);
    });
  } else if(_fitOrder == 9) {
    ddfor(emGrid, [&](const iVec &idx) {
      outData(idx) =
          (-aData(idx + u * 4) * 3 + aData(idx + u * 3) * 32 - aData(idx + u * 2) * 168 + aData(idx + u) * 672
              - aData(idx - u) * 672 + aData(idx - u * 2) * 168 - aData(idx - u * 3) * 32 + aData(idx - u * 4) * 3)
              / (840.0 * dx[pD]);
    });
  } else {
    assert(0);
  }
  // 2. apply the correction.
  matOp.AXPY(1.0, aData, [](iVec) { return 0.0; }, outData);
}

//============================================================
#if DIM == 2
template class PartialDer<2, 0, 5>;
template class PartialDer<2, 1, 5>;
template class PartialDer<2, 0, 7>;
template class PartialDer<2, 1, 7>;
template class PartialDer<2, 0, 9>;
template class PartialDer<2, 1, 9>;
#elif DIM == 3
template class PartialDer<3, 0>;
template class PartialDer<3, 1>;
template class PartialDer<3, 2>;
#endif
